#!/usr/bin/env python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import enum
from ._main import WinePrefix


class WineRegistryEditor:

    class ValueType(enum.Enum):                     # from https://docs.microsoft.com/en-us/windows/win32/sysinfo/registry-value-types
        REG_NONE = enum.auto()
        REG_DWORD = enum.auto()
        REG_DWORD_LITTLE_ENDIAN = enum.auto()
        REG_DWORD_BIG_ENDIAN = enum.auto()
        REG_QWORD = enum.auto()
        REG_QWORD_LITTLE_ENDIAN = enum.auto()
        REG_SZ = enum.auto()
        REG_EXPAND_SZ = enum.auto()
        REG_MULTI_SZ = enum.auto()
        REG_BINARY = enum.auto()
        REG_LINK = enum.auto()

    class ParseError(Exception):
        pass

    class ValueTypeError(Exception):
        pass

    def __init__(self, parent):
        '''
        This object represents the functions of regedit.
        Note: regedit may or may not be executed.

        Parameters
        ----------
        parent: WinePrefix
            parent object
        '''

        self._parent = parent

        self._files = {
            "HKEY_LOCAL_MACHINE": _RegFile(os.path.join(self._parent._path, "system.reg"), relative_to="\\\\Machine"),
            "HKEY_CURRENT_USER": _RegFile(os.path.join(self._parent._path, "user.reg"), relative_to="\\\\User\\\\S-1-5-21-0-0-0-1000"),
        }
        self._changed = {k: False for k in self._files}

    def is_key_exist(self, key):
        ret = self._tryGet(key)
        if ret is not None:
            key, fileObj = ret
            return fileObj.is_key_exist(key)
        else:
            return False

    def create_key(self, key):
        key, fileObj = self._get(key)
        fileObj.create_key(key)

    def list_all_subkeys(self, key):
        key, fileObj = self._get(key)
        return fileObj.list_subkeys(key)

    def delete_all_subkeys(self, key):
        key, fileObj = self._get(key)
        for k in fileObj.list_subkeys(key):
            fileObj.delete_key(k)

    def delete_key(self, key, recursive=False):
        key, fileObj = self._get(key)
        if not fileObj.is_key_exist(key):
            return
        fileObj.delete_key(key, recursive=recursive)

    def list_values(self, key):
        key, fileObj = self._get(key)
        return fileObj.get_key(key).get_value_names()

    def is_value_exist(self, key, value_name, value_type):
        key, fileObj = self._get(key)
        keyObj = fileObj.get_key(key)
        if keyObj.is_value_exist(value_name):
            if keyObj.get_value(value_name)[0] != value_type:
                raise self.ValueTypeError("invalid value type")          # FIXME: optimize error message
            return True
        else:
            return False

    def get_value(self, key, value_name, value_type, fallback=None):
        key, fileObj = self._get(key)
        keyObj = fileObj.get_key(key)
        if fallback is not None and not keyObj.is_value_exist(value_name):
            return fallback
        valueType, data = keyObj.get_value(value_name)
        if valueType != value_type:
            raise self.ValueTypeError("invalid value type")              # FIXME: optimize error message
        return data

    def set_value(self, key, value_name, value_type, data):
        key, fileObj = self._get(key)
        fileObj.get_key(key).set_value(value_name, value_type, data)

    def delete_value(self, key, value_name):
        key, fileObj = self._get(key)
        fileObj.get_key(key).delete_value(value_name)

    def delete_all_values(self, key):
        key, fileObj = self._get(key)
        keyObj = fileObj.get_key(key)
        for value_name in keyObj.get_value_names():
            keyObj.delete_value(value_name)

    def import_file(self, path):
        obj = self._RegImportTemplateFile(path)
        for name, keyObj in obj.keys.items():
            if not self.is_key_exist(name):
                self.create_key(name)
            for value_name in keyObj.get_value_names():
                self.set_value(name, value_name, *keyObj.get_value(value_name))

    def save(self):
        for k, v in self._files.items():
            v.save()
            # if self._changed[k]:
            #     v.save()

    def _tryGet(self, key):
        for prefix, fileObj in self._files.items():
            prefix += "\\"
            if key.startswith(prefix):
                return (key[len(prefix):], fileObj)
        return None

    def _get(self, key):
        ret = self._tryGet(key)
        assert ret is not None
        return ret


class _RegFile:

    def __init__(self, path, relative_to):
        assert relative_to is not None

        self._path = path
        self._version = None
        self._relative_to = None
        self._arch = None
        self._keys = dict()

        registry_lines = None
        with open(self._path, "r", encoding='utf-8') as f:
            registry_lines = [x.rstrip("\n") for x in f.readlines()]

        temp_lines = []
        lineNo = 0
        for i in range(0, len(registry_lines)):
            lineNo = i + 1
            line = registry_lines[i].rstrip()
            if line == "":
                pass
            elif line.startswith(self._version_header):                       # example: "WINE REGISTRY Version 2"
                if self._version is not None:
                    raise WineRegistryEditor.ParseError("duplicate version header line %d \"%s\" in registry file \"%s\"" % (lineNo, registry_lines[i], self._path))
                self._version = int(line[len(self._version_header):])
            elif line.startswith(self._relative_to_header):                   # example: ";; All keys relative to \\User\\S-1-5-21-0-0-0-1000"
                if self._relative_to is not None:
                    raise WineRegistryEditor.ParseError("duplicate relative-to header line %d \"%s\" in registry file \"%s\"" % (lineNo, registry_lines[i], self._path))
                r = line[len(self._relative_to_header):]
                if r != relative_to:
                    raise WineRegistryEditor.ParseError("invalid relative-to header line %d \"%s\" in registry file \"%s\"" % (lineNo, registry_lines[i], self._path))
                self._relative_to = r
            elif line.startswith("#arch"):                                    # example: "#arch=win64"
                self._arch = line.split("=")[1]
            elif line.startswith("["):                                        # example: "[Software\\Classes\\.ai]"
                if len(temp_lines) > 0:
                    k = _RegKey(filePath=self._path, lineNo=(lineNo - len(temp_lines)), lines=temp_lines)
                    self._keys[k.name] = k
                temp_lines = [line]
            elif len(temp_lines) > 0:
                temp_lines.append(line)                                     # variable temp_lines includes empty lines
            else:
                raise WineRegistryEditor.ParseError("invalid line %d \"%s\" in registry file \"%s\"" % (lineNo, registry_lines[i], self._path))

        if len(temp_lines) > 0:
            k = _RegKey(filePath=self._path, lineNo=(lineNo - len(temp_lines)), lines=temp_lines)
            self._keys[k.name] = k

        if self._version is None:
            raise WineRegistryEditor.ParseError("no version header line in registry file \"%s\"" % (self._path))
        if self._relative_to is None:
            raise WineRegistryEditor.ParseError("no relative-to header line in registry file \"%s\"" % (self._path))
        if self._arch is None:
            raise WineRegistryEditor.ParseError("no arch information in registry file \"%s\"" % (self._path))

    @property
    def version(self):
        return self._version

    @property
    def arch(self):
        return self._arch

    def is_key_exist(self, key):
        assert key is not None
        return key in self._keys

    def get_key(self, key):
        assert key is not None
        return self._keys[key]

    def list_subkeys(self, key, recursive=False):
        assert key is not None and key in self._keys
        ret = list(self._keys)
        ret = [x for x in ret if x.startswith(key + "\\")]
        if not recursive:
            ret = [x for x in ret if "\\" not in x[len(key)+1:]]
        ret.sort()
        return ret

    def create_key(self, key):
        assert key is not None and key not in self._keys
        # FIXME: assert parent key exists
        k = _RegKey(name=key)
        self._keys[k.name] = k
        return k

    def delete_key(self, key, recursive=False):
        if key not in self._keys:
            return
        tlist = self.list_subkeys(key, recursive=True)
        if not recursive:
            assert len(tlist) == 0
        for k in reversed(tlist):
            self.delete_key(k, recursive=False)
        del self._keys[key]

    def render(self):
        content = "{}{}\n".format(self._version_header, self._version)
        content += "{}{}\n\n".format(self._relative_to_header, self._relative_to)
        content += "#arch={}\n".format(self._arch)
        for key in self._keys:
            content += "\n"
            content += self._keys[key].render()
        return content

    def save(self):
        with open(self._path, "w", encoding='utf-8') as f:
            f.write(self.render())

    _version_header = "WINE REGISTRY Version "
    _relative_to_header = ";; All keys relative to "


class _RegKey:

    def __init__(self, filePath=None, lineNo=None, lines=None, name=None):
        if filePath is not None:
            assert lineNo is not None and lines is not None
            assert name is None
        else:
            assert lineNo is None and lines is None
            assert name is not None

        self._name = None
        self._rawName = None
        self._timestamp = None
        self._rawTimestamp = None
        self._metas = dict()
        self._rawValues = dict()        # we keep raw values so that we can render them back

        def __rawName2name(rawName):
            return rawName.replace("\\\\", "/").strip("[]")

        def __name2rawName(name):
            return "[{}]".format(name.replace("/", "\\\\"))

        def __rawTimestamp2timestamp(rawTimestamp):
            # Parse timestamp either as int or float
            ts_parts = rawTimestamp.strip().split()
            if len(ts_parts) == 1:
                return int(ts_parts[0])
            else:
                return float("{}.{}".format(ts_parts[0], ts_parts[1]))

        if filePath is None:
            # Key is created by path, it's a new key
            timestamp = datetime.now().timestamp()
            self._name = name
            self._rawName = __name2rawName(self._name)
            self._rawTimestamp = " ".join(str(timestamp).split("."))
            self._timestamp = __rawTimestamp2timestamp(self._rawTimestamp)
            self._metas["time"] = WindowsFileTime.from_unix_timestamp(timestamp).to_hex()
        else:
            # Existing key loaded from file
            # lines example:
            #   [AppEvents\\Schemes\\Apps\\Explorer\\Navigating\\.Current] 1649505057
            #   #time=1d84c08138e6746
            #   @=""
            #   "ColorTable00"=dword:00000000
            #   "CaptionFont"=hex:f5,ff,ff,ff,00,00,00,00,00,00,00,00,00,00,00,00,bc,02,00,00,\
            #     00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00
            #   "BorderWidth"="1"
            self._rawName, self._rawTimestamp = re.split(re.compile(r"(?<=[^\\]\]) "), lines[0], maxsplit=1)
            self._name = __rawName2name(self._rawName)
            self._timestamp = __rawTimestamp2timestamp(self._rawTimestamp)

            line = ""
            for cl in lines[1:]:
                line += cl
                if line.endswith("\\"):
                    line += "\n"
                    continue

                if line.startswith("#"):
                    if "\n" in line:
                        raise WineRegistryEditor.ParseError("invalid line %d \"%s\" in registry file \"%s\"" % (lineNo, line, filePath))
                    parts = line[1:].split("=")
                    if len(parts) == 2:
                        name, value = parts
                    elif len(parts) == 1:
                        name, value = parts[0], None
                    else:
                        raise WineRegistryEditor.ParseError("invalid line %d \"%s\" in registry file \"%s\"" % (lineNo, line, filePath))
                    if name in self._metas:
                        raise WineRegistryEditor.ParseError("duplicate meta \"%s\" in registry file \"%s\"" % (name, filePath))
                    self._metas[name] = value
                elif line.startswith('"'):
                    try:
                        name, value = re.split(re.compile(r"(?<![^\\]\\\")="), line, maxsplit=1)
                    except ValueError:
                        raise WineRegistryEditor.ParseError("invalid line %d \"%s\" in registry file \"%s\"" % (lineNo, line, filePath))
                    if "\n" in name:
                        raise WineRegistryEditor.ParseError("invalid line %d \"%s\" in registry file \"%s\"" % (lineNo, line, filePath))
                    name = name[1:-1]               # remove leading quote and trailing quote
                    self._rawValues[name] = value
                elif line.startswith("@"):
                    self._rawValues["default"] = line.split("=", 1)[1]
                elif line == "":
                    pass
                else:
                    raise WineRegistryEditor.ParseError("invalid line %d \"%s\" in registry file \"%s\"" % (lineNo, line, filePath))

                line = ""

    @property
    def name(self):
        return self._name

    @property
    def raw_name(self):
        return self._rawName

    @property
    def timestamp(self):
        return self._timestamp

    @property
    def raw_timestamp(self):
        return self._rawTimestamp

    def has_metadata(self, metadata_name):
        return metadata_name in self._metas

    def get_metadata(self, metadata_name):
        return self._metas[metadata_name]

    def get_value_names(self):
        return list(self._rawValues.keys())

    def is_value_exist(self, value_name):
        return value_name in self._rawValues

    def get_value(self, value_name):
        rawValue = self._rawValues[value_name]
        assert "\n" not in rawValue     # FIXME: we don't support multi-line currently

        if rawValue.startswith('"') and rawValue.endswith('"'):
            return (WineRegistryEditor.ValueType.REG_SZ, Util.decodeUnicode(rawValue[1:-1]))

        if rawValue.startswith("dword:"):
            return (WineRegistryEditor.ValueType.REG_DWORD, int(rawValue[6:], 16))

        raise NotImplementedError("TODO")

    def get_raw_value(self, value_name):
        return self._rawValues[value_name]

    def set_value(self, value_name, value_type, data):
        if value_type == WineRegistryEditor.ValueType.REG_DWORD:
            rawValue = "dword:{:08x}".format(data)
        elif value_type == WineRegistryEditor.ValueType.REG_SZ:
            rawValue = '"{}"'.format(Util.encodeUnicode(data))
        else:
            raise NotImplementedError("TODO")

        self._rawValues[value_name] = rawValue

    def delete_value(self, value_name):
        del self._rawValues[value_name]

    def render(self):
        """Return the content of the key in the wine .reg format"""

        content = self._rawName + " " + self._rawTimestamp + "\n"
        for key in sorted(self._metas):
            value = self._metas[key]
            if value is None:
                content += "#{}\n".format(key)
            else:
                content += "#{}={}\n".format(key, value)
        if "default" in self._rawValues:
            content += "@={}\n".format(self._rawValues["default"])
        for key in sorted(self._rawValues):
            if key == "default":
                continue
            content += "\"{}\"={}\n".format(key, self._rawValues[key])
        return content


class _RegImportTemplateFile:

    def __init__(self, path):
        self._keys = dict()

        registry_lines = None
        with open(path, "r", encoding='utf-8') as f:
            registry_lines = [x.rstrip("\n") for x in f.readlines()]

        temp_lines = []
        lineNo = 0
        for i in range(0, len(registry_lines)):
            lineNo = i + 1
            line = registry_lines[i].rstrip()

            if line == "":
                pass
            elif line.startswith("["):                                        # example: "[Software\\Classes\\.ai]"
                if len(temp_lines) > 0:
                    k = _RegImportTemplateKey(path, lineNo - len(temp_lines), temp_lines)
                    self._keys[k.name] = k
                temp_lines = [line]
            elif len(temp_lines) > 0:
                temp_lines.append(line)                                     # variable temp_lines includes empty lines
            else:
                raise WineRegistryEditor.ParseError("invalid line %d \"%s\" in registry template file \"%s\"" % (lineNo, registry_lines[i], path))

        if len(temp_lines) > 0:
            k = _RegImportTemplateKey(path, lineNo - len(temp_lines), temp_lines)
            self._keys[k.name] = k

    @property
    def keys(self):
        return self._keys


class _RegImportTemplateKey:

    def __init__(self, filePath, lineNo, lines):
        # lines example (note: different from wine registry file):
        #   [AppEvents\Schemes\Apps\Explorer\Navigating\.Current]
        #   @=""
        #   "ColorTable00"=dword:00000000
        #   "CaptionFont"=hex:f5,ff,ff,ff,00,00,00,00,00,00,00,00,00,00,00,00,bc,02,00,00,\
        #     00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00
        #   "BorderWidth"="1"
        self._name = re.fullmatch(re.compile(r"\[(.*)\]"), lines[0]).group(0).replace("\\", "/")

        self._rawValues = dict()        # we keep raw values so that we can render them back
        line = ""
        for cl in lines[1:]:
            line += cl
            if line.endswith("\\"):
                line += "\n"
                continue

            if line == "":
                pass
            elif line.startswith('"'):
                try:
                    name, value = re.split(re.compile(r"(?<![^\\]\\\")="), line, maxsplit=1)
                except ValueError:
                    raise WineRegistryEditor.ParseError("invalid line %d \"%s\" in file \"%s\"" % (lineNo, line, filePath))
                if "\n" in name:
                    raise WineRegistryEditor.ParseError("invalid line %d \"%s\" in file \"%s\"" % (lineNo, line, filePath))
                name = name[1:-1]               # remove leading quote and trailing quote
                self._rawValues[name] = value
                line = ""
            elif line.startswith("@"):
                self._rawValues["default"] = line.split("=", 1)[1]
                line = ""
            else:
                raise WineRegistryEditor.ParseError("invalid line %d \"%s\" in file \"%s\"" % (lineNo, line, filePath))

    @property
    def name(self):
        return self._name

    def get_value_names(self):
        return list(self._rawValues.keys())

    def get_value(self, value_name):
        rawValue = self._rawValues[value_name]
        assert "\n" not in rawValue     # FIXME: we don't support multi-line currently

        if rawValue.startswith('"') and rawValue.endswith('"'):
            return (WineRegistryEditor.ValueType.REG_SZ, Util.decodeUnicode(rawValue[1:-1]))

        if rawValue.startswith("dword:"):
            return (WineRegistryEditor.ValueType.REG_DWORD, int(rawValue[6:], 16))

        raise NotImplementedError("TODO")
